

#include "ednn_maxpool.h"
#include "ednn_backends.h"


static ednn_ret maxpool_forward(struct ednn_layer *layer)
{
    ednn_uint16_t out_h, out_w;
    struct ednn_layer_maxpool *maxpool = (struct ednn_layer_maxpool *)layer;
    struct ednn_layer *prev = ednn_list_entry(layer->list.prev, 
        struct ednn_layer, list);

    if (layer->output->ndim == 1) { // global pooling
        out_h = 1; 
        out_w = 1;
    } else {
        out_h = layer->output->dim[0];
        out_w = layer->output->dim[1];
    }

    local_maxpool_q7_hwc_nonsquare(
        prev->output->pd,
        prev->output->dim[0],
        prev->output->dim[1],
        prev->output->dim[2],
        maxpool->kernel.h,
        maxpool->kernel.w,
        maxpool->padding.h,
        maxpool->padding.w,
        maxpool->stride.h,
        maxpool->stride.w,
        layer->output->pd,
        out_h, 
        out_w
    );

    return ednn_ok;
}

static ednn_ret maxpool_build(struct ednn_layer *layer)
{
    struct ednn_layer_maxpool *maxpool = (struct ednn_layer_maxpool *)layer;
    struct ednn_layer *prev = ednn_list_entry(layer->list.prev, 
        struct ednn_layer, list);
    
    ednn_shape_t dim[3];
    if (maxpool->padding_type == EDNN_SAME_PADDING) {
        dim[0] = ednn_ceilif(prev->output->dim[0], maxpool->stride.h);
        dim[1] = ednn_ceilif(prev->output->dim[1], maxpool->stride.w);
        dim[2] = prev->output->dim[2];
    } else {
        dim[0] = ednn_ceilif(prev->output->dim[0] - maxpool->kernel.h + 1, 
            maxpool->stride.h);
        dim[1] = ednn_ceilif(prev->output->dim[1] - maxpool->kernel.w + 1, 
            maxpool->stride.w);
        dim[2] = prev->output->dim[2];
    }

    layer->output = ednn_tensor_create(dim, 3);

    return ednn_ok;
}

struct ednn_layer *ednn_maxpool_layer(ednn_shape3d_t kernel, 
    ednn_shape3d_t stride, ednn_padding_type padding_type)
{
    struct ednn_layer_maxpool *layer = ednn_mem_zalloc(sizeof(struct ednn_layer_maxpool));
    ednn_assert(layer != EDNN_NULL);

    layer->super.type    = EDNN_LAYER_MAXPOOL;
    layer->super.build   = maxpool_build;
    layer->super.forward = maxpool_forward;

    layer->kernel       = kernel;
    layer->stride       = stride;
    layer->padding_type = padding_type;

    if (layer->padding_type == EDNN_SAME_PADDING) {
        layer->padding.h = (kernel.h - 1) / 2;
        layer->padding.w = (kernel.w - 1) / 2;
        layer->padding.c = 1;
    } else {
        layer->padding.h = 0;
		layer->padding.w = 0;
		layer->padding.c = 0;
    }

    return (struct ednn_layer *)layer;
}